import gradio as gr
import cv2 as cv
import matplotlib.pyplot as plt
import numpy as np
from pinecone import Pinecone
import logging
import time

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 初始化Pinecone
pc = Pinecone(api_key="c64273d1-3f4a-4ae3-99a1-4a0a0b7a5f87")
index_name = "mnist-index"
index = pc.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 定义图像处理函数
def preprocess_image(image):
    if image is None:
        return None
    image = cv.cvtColor(image, cv.COLOR_RGBA2GRAY)
    plt.imsave("preprocessed_image.png", image, cmap='gray')
    image = cv.resize(image, (28, 28), interpolation=cv.INTER_AREA)
    plt.imsave("resized_image.png", image, cmap="gray")
    epsilon = 1e-8
    image = image.astype('float32') / 255.0 + epsilon
    return image.ravel()

# 定义查询Pinecone的函数
def query_pinecone(image_vector, max_retries=3):
    for attempt in range(max_retries):
        try:
            query_result = index.query(
                vector=image_vector,
                top_k=11,
                include_metadata=True
            )
            return query_result
        except Exception as e:
            wait_time = 2 ** attempt
            logging.warning(f"查询失败，等待{wait_time}秒后重试（第{attempt+1}次）: {str(e)}")
            time.sleep(wait_time)
            if attempt == max_retries - 1:
                logging.error(f"查询失败，已达到最大重试次数: {str(e)}")
                return None

# 定义预测函数
def predict(image):
    image_vector = preprocess_image(image)
    if image_vector is None:
        return "图像为空"
    query_result = query_pinecone(image_vector)
    if query_result is None:
        return "查询失败"
    nearest_labels = [match['metadata']['label'] for match in query_result['matches']]
    predicted_label = max(set(nearest_labels), key=nearest_labels.count)
    return str(predicted_label)

# 构建Gradio界面
gr.Interface(fn=predict, inputs=gr.ImageEditor(label="上传图像"), outputs=gr.Label(num_top_classes=1)).launch()